from typing import Optional

import moderngl
import numpy as np

from .attribute import Attribute, Constant


class Geometry:
    attributes: dict[Constant, Attribute] = {
        Constant.POSITION: Attribute('position', '3f'),
    }
    mode = moderngl.TRIANGLES

    def __init__(self):
        super().__init__()
        self.position: Optional[np.ndarray] = None


class Plane(Geometry):
    attributes: dict[Attribute, str] = {
        Constant.POSITION: Attribute('position', '3f'),
        Constant.NORMAL: Attribute('normal', '3f'),
        Constant.TEXCOORD: Attribute('uv', '2f'),
        Constant.INDEX: Attribute('index', 'i'),
    }

    def __init__(
            self,
            size: tuple[float, float] = (1.0, 1.0),
            segments: tuple[int, int] = (1, 1)
    ):
        super().__init__()
        width, height = size
        half_w, half_h = width / 2, height / 2
        sw, sh = segments
        sw1, sh1 = sw + 1, sh + 1
        dw, dh = width / sw, height / sh

        points = sw1 * sh1
        triangles = sw * sh * 2
        positions = np.ndarray((points, 3), dtype='f4')
        uvs = np.ndarray((points, 2), dtype='f4')
        normals = np.ndarray((points, 3), dtype='f4')
        index = np.ndarray((triangles, 3), dtype='i4')

        i = 0
        for ix in range(sw1):
            for iz in range(sh1):
                x = ix * dw - half_w
                z = iz * dh - half_h
                y = 0
                u = ix / sw
                v = iz / sh
                positions[i] = x, y, z
                uvs[i] = u, v
                normals[i] = 0, 1, 0
                i += 1

        i = 0
        for ix in range(sw):
            tmp = ix * sh1
            for iz in range(sh):
                a = tmp + iz
                b = a + 1
                c = b + sh1
                d = c - 1
                index[i] = a, b, c
                index[i + 1] = d, a, c
                i += 2

        self.position = positions
        self.uv = uvs
        self.normal = normals
        self.index = index
